from re import S
from datasets.basedataset import BaseDataset
from nets.bases.train_base import TrainBase
from nets.bases.basicmodel import BasicModule
from datasets.martian_dataset import MartianDataSet
from nets.models.unet_model import UNetModel, DepthUNetModel
import torch as t
import torch.nn as nn
import torch.nn.functional as F


def dice_loss(input, target):
    smooth = 1e-6
    iflat = input.view(-1)
    tflat = target.view(-1)
    intersection = (iflat * tflat).sum()
    return ((2. * intersection + smooth) / (iflat.sum() + tflat.sum() + smooth))

def mul_dice_loss(input, target):
    loss = 0
    for i in range(input.shape[0]):
        loss += dice_loss(input[i], target[i])
    return loss / input.shape[0]

def batch_multiclass_dice_loss(input, target):
    loss = 0
    for i in range(input.shape[0]):
        loss += mul_dice_loss(input[i], target[i])
    return loss / input.shape[0]

class UNetTrain(TrainBase):
    
    def __init__(self, args):
        if args.isTranning == True:
            self.create_train(MartianDataSet(r'I:\Datasets\工业检测\样本1\Image1_B_result'), DepthUNetModel(), 
                            batch_size=args.batch_size, shuffle=args.shuffle, num_workers=args.num_workers, size=args.size, 
                            optim=args.optim, 
                            lr=args.lr, save_model_count=args.save_model_count, device=args.device, 
                            is_print_model_info=args.is_print_model_info)  # 初始化父类
            self.cross_loss = nn.BCELoss()
            
        else:
            self.create_test(DepthUNetModel(), args.size, args.device, args.model_path),
        
    
    
    def train_once(self, datas):
        img, mask_img = datas
        img = img.to(self.device)
        mask_img = mask_img.to(self.device)
        #计算
        res = self.model(img)
        return [res, mask_img]
    
    
    def caculator_loss(self, res):
        res_img, mask_img = res
        # print(res_img.shape, mask_img.shape)
        # res_img = t.squeeze(res_img)
        loss1 = self.cross_loss(t.sigmoid(res_img), mask_img)
        loss2 = batch_multiclass_dice_loss(F.softmax(res_img, dim=1), mask_img)
        return loss1 + loss2    
        
        
    def get_model_test_size(self):
        return (3, 512, 512)
    

        